"""
SimSwap GAN Simulator模块
功能：使用SimSwap模型进行人脸交换，替代原有的Faceswap
这是FakeTagger系统的第3个模块
"""

import torch
import torch.nn as nn
import os
import sys
from PIL import Image
import numpy as np
import cv2
import warnings

class SimSwapGANSimulator:
    """
    SimSwap GAN模拟器
    使用SimSwap模型替代原有的Faceswap进行深度人脸交换
    支持您训练过的SimSwap权重
    """
    
    def __init__(self, device='cuda', simswap_model_path=None, use_pretrained_weights=True):
        """
        初始化SimSwap GAN模拟器
        
        Args:
            device: 设备类型
            simswap_model_path: SimSwap模型路径（您训练过的权重）
            use_pretrained_weights: 是否使用预训练权重
        """
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        self.model = None
        self.is_available = False
        self.use_simulation = False
        self.simswap_model_path = simswap_model_path
        
        print(f"SimSwap GAN Simulator initializing...")
        print(f"  - Device: {self.device}")
        print(f"  - Model path: {simswap_model_path}")
        print(f"  - Use pretrained weights: {use_pretrained_weights}")
        
        # 尝试加载SimSwap模型
        self._load_simswap_model()
    
    def _load_simswap_model(self):
        """加载SimSwap模型"""
        try:
            # 检查SimSwap项目路径
            simswap_paths = [
                './SimSwap-main',
                '../SimSwap-main',
                './SimSwap',
                '../SimSwap',
                os.path.join(os.path.dirname(__file__), 'SimSwap-main'),
                os.path.join(os.path.dirname(__file__), '..', 'SimSwap-main')
            ]
            
            simswap_path = None
            for path in simswap_paths:
                if os.path.exists(path):
                    simswap_path = path
                    break
            
            if simswap_path is None:
                raise ImportError("SimSwap directory not found")
            
            print(f"  - Found SimSwap at: {simswap_path}")
            
            # 获取绝对路径
            simswap_abs_path = os.path.abspath(simswap_path)
            print(f"  - Absolute path: {simswap_abs_path}")
            
            # 添加SimSwap路径到系统路径
            if simswap_abs_path not in sys.path:
                sys.path.insert(0, simswap_abs_path)
                print(f"  - Added to sys.path: {simswap_abs_path}")
            
            # 保存当前工作目录
            original_cwd = os.getcwd()
            
            try:
                # 切换到SimSwap目录（确保相对路径正确）
                os.chdir(simswap_abs_path)
                print(f"  - Changed working directory to: {simswap_abs_path}")
                
                # 尝试导入SimSwap模块
                print(f"  - Trying to import modules...")
                from models.models import create_model
                from options.test_options import TestOptions
                print(f"  - SimSwap modules imported successfully")
                
                # 立即恢复工作目录
                os.chdir(original_cwd)
                
            except ImportError as e:
                # 恢复工作目录
                os.chdir(original_cwd)
                # 如果导入失败，尝试其他方法
                print(f"  - Standard import failed: {e}")
                print(f"  - Trying alternative import method...")
                
                # 方法2：使用importlib动态导入
                import importlib.util
                
                # 导入models.models
                models_path = os.path.join(simswap_abs_path, 'models', 'models.py')
                if os.path.exists(models_path):
                    spec = importlib.util.spec_from_file_location("models.models", models_path)
                    models_module = importlib.util.module_from_spec(spec)
                    sys.modules["models.models"] = models_module
                    spec.loader.exec_module(models_module)
                    create_model = models_module.create_model
                    print(f"  - models.models imported via importlib")
                else:
                    raise ImportError(f"models.models.py not found at {models_path}")
                
                # 导入options.test_options  
                options_path = os.path.join(simswap_abs_path, 'options', 'test_options.py')
                if os.path.exists(options_path):
                    spec = importlib.util.spec_from_file_location("options.test_options", options_path)
                    options_module = importlib.util.module_from_spec(spec)
                    sys.modules["options.test_options"] = options_module
                    spec.loader.exec_module(options_module)
                    TestOptions = options_module.TestOptions
                    print(f"  - options.test_options imported via importlib")
                else:
                    raise ImportError(f"test_options.py not found at {options_path}")
                    
            except Exception as e:
                # 确保恢复工作目录
                os.chdir(original_cwd)
                raise ImportError(f"Failed to import SimSwap modules: {e}")
            
            # 创建测试选项
            opt = TestOptions()
            opt.initialize()
            
            # 构建绝对路径
            arcface_path = os.path.join(simswap_abs_path, 'arcface_model', 'arcface_checkpoint.tar')
            checkpoints_dir = os.path.join(simswap_abs_path, 'checkpoints')
            
            # 检查可用的模型目录
            available_models = []
            for model_dir in ['simswap512', 'people', 'simswap256', 'simswap128']:
                model_path = os.path.join(checkpoints_dir, model_dir)
                if os.path.exists(model_path):
                    available_models.append(model_dir)
            
            # 选择最佳模型（优先级：simswap512 > people > others）
            if 'simswap512' in available_models:
                model_name = 'simswap512'
                print(f"  - Using SimSwap512 model")
            elif 'people' in available_models:
                model_name = 'people'
                print(f"  - Using People model")
            elif available_models:
                model_name = available_models[0]
                print(f"  - Using {model_name} model")
            else:
                raise ImportError(f"No SimSwap models found in {checkpoints_dir}")
            
            # 设置基本参数
            opt.parser.set_defaults(
                name=model_name,
                Arc_path=arcface_path,
                pic_a_path=os.path.join(simswap_abs_path, 'demo_file/Iron_man.jpg'),
                pic_b_path=os.path.join(simswap_abs_path, 'demo_file/multi_people.jpg'),
                pic_specific_path=os.path.join(simswap_abs_path, 'demo_file/specific_person.jpg'),
                multisepcific_dir=os.path.join(simswap_abs_path, 'demo_file/multispecific'),
                video_path=os.path.join(simswap_abs_path, 'demo_file/multi_people.mp4'),
                temp_path=os.path.join(simswap_abs_path, 'temp_results'),
                output_path=os.path.join(simswap_abs_path, 'output/'),
                checkpoints_dir=checkpoints_dir,
                result_dir=os.path.join(simswap_abs_path, 'results'),
                id_thres=0.03,
                no_simswaplogo=False,
                use_mask=False,
                crop_size=224
            )
            
            # 如果提供了自定义模型路径，使用它
            if self.simswap_model_path:
                # 确保路径是绝对路径
                if not os.path.isabs(self.simswap_model_path):
                    custom_model_path = os.path.abspath(self.simswap_model_path)
                else:
                    custom_model_path = self.simswap_model_path
                    
                if os.path.exists(custom_model_path):
                    opt.parser.set_defaults(
                        checkpoints_dir=os.path.dirname(custom_model_path),
                        name=os.path.basename(custom_model_path).split('.')[0]
                    )
                    print(f"  - Using custom model: {custom_model_path}")
                else:
                    print(f"  - Warning: Custom model path not found: {custom_model_path}")
                    print(f"  - Falling back to default model path")
            
            self.opt = opt.parse()
            
            # 检查关键路径是否存在
            print(f"  - Checking arcface path: {self.opt.Arc_path}")
            if not os.path.exists(self.opt.Arc_path):
                print(f"  - Warning: Arcface model not found at {self.opt.Arc_path}")
            
            print(f"  - Checking checkpoints dir: {self.opt.checkpoints_dir}")
            if not os.path.exists(self.opt.checkpoints_dir):
                print(f"  - Creating checkpoints directory: {self.opt.checkpoints_dir}")
                os.makedirs(self.opt.checkpoints_dir, exist_ok=True)
            
            # 创建模型
            print(f"  - Creating SimSwap model...")
            
            # 临时禁用PyTorch 2.6的安全权重加载限制
            import torch
            original_load = torch.load
            def safe_load(*args, **kwargs):
                kwargs['weights_only'] = False  # 允许加载自定义类
                return original_load(*args, **kwargs)
            
            # 临时替换torch.load
            torch.load = safe_load
            
            try:
                self.model = create_model(self.opt)
                print(f"  - Model created successfully")
            finally:
                # 恢复原始的torch.load
                torch.load = original_load
            
            self.model.eval()
            print(f"  - Model set to evaluation mode")
            
            # 冻结所有参数（用于推理）
            param_count = 0
            for param in self.model.parameters():
                param.requires_grad = False
                param_count += param.numel()
                
            self.is_available = True
            print(f"  - SimSwap model loaded successfully")
            print(f"  - Model parameters frozen for inference ({param_count:,} parameters)")
            
        except Exception as e:
            print(f"  - SimSwap loading failed: {e}")
            print(f"  - Switching to simulation mode")
            self.use_simulation = True
            self.is_available = True  # 模拟模式也算可用
    
    def deepfake(self, encoded_images, swap_type='A'):
        """
        进行深度人脸交换
        
        Args:
            encoded_images: 嵌入水印的图像 (batch_size, 3, H, W)，范围[-1,1]或[0,1]
            swap_type: 交换类型
                      'A' 或 'trump_to_cage': Trump -> Cage
                      'B' 或 'cage_to_trump': Cage -> Trump
                      
        Returns:
            torch.Tensor: 深度伪造后的图像，与输入相同格式
        """
        if self.use_simulation:
            return self._simulate_face_swap(encoded_images, swap_type)
        
        try:
            # 确保输入在正确的设备上
            encoded_images = encoded_images.to(self.device)
            
            # 记住原始输入范围
            original_is_normalized = encoded_images.min() < 0  # 是否是[-1,1]范围
            
            # 确保输入范围正确（SimSwap期望[0,1]）
            if original_is_normalized:  # 如果输入是[-1,1]
                encoded_images = (encoded_images + 1.0) / 2.0  # 转换到[0,1]
            
            batch_size = encoded_images.shape[0]
            
            # 为每种交换类型创建不同的身份特征
            if not hasattr(self, 'identity_features'):
                self._initialize_identity_features()
            
            with torch.no_grad():
                # 选择身份特征
                if swap_type in ['A', 'trump_to_cage']:
                    # Trump -> Cage: 使用Cage的身份特征
                    latent_id = self.identity_features['cage'].repeat(batch_size, 1)
                elif swap_type in ['B', 'cage_to_trump']:
                    # Cage -> Trump: 使用Trump的身份特征
                    latent_id = self.identity_features['trump'].repeat(batch_size, 1)
                else:
                    raise ValueError(f"Unknown swap_type: {swap_type}")
                
                # 调用SimSwap模型 (正确的接口)
                # model(img_id, img_att, latent_id, latent_att, for_G)
                deepfaked = self.model(None, encoded_images, latent_id, None, True)
                
                # SimSwap在推理模式下只返回图像，不返回损失
                if isinstance(deepfaked, list):
                    deepfaked = deepfaked[0] if len(deepfaked) > 0 else deepfaked
                
                # 确保输出范围正确，转换回输入的原始格式
                deepfaked = torch.clamp(deepfaked, 0, 1)
                
                # 如果原始输入是[-1,1]范围，转换回去
                if original_is_normalized:
                    deepfaked = deepfaked * 2.0 - 1.0
                
                return deepfaked
                
        except Exception as e:
            print(f"  - SimSwap inference failed: {e}")
            print(f"  - Falling back to simulation mode")
            return self._simulate_face_swap(encoded_images, swap_type)
    
    def _initialize_identity_features(self):
        """初始化身份特征（Trump和Cage的平均特征）"""
        print(f"  - Initializing identity features...")
        
        # 创建虚拟的身份特征向量
        # 在实际应用中，这些应该是从真实的Trump和Cage图像中提取的
        feature_dim = 512  # ArcFace特征维度
        
        # 创建不同的身份特征
        torch.manual_seed(42)  # 确保可重复性
        trump_feature = torch.randn(1, feature_dim, device=self.device)
        trump_feature = torch.nn.functional.normalize(trump_feature, p=2, dim=1)
        
        torch.manual_seed(123)  # 不同的种子产生不同特征
        cage_feature = torch.randn(1, feature_dim, device=self.device)
        cage_feature = torch.nn.functional.normalize(cage_feature, p=2, dim=1)
        
        self.identity_features = {
            'trump': trump_feature,
            'cage': cage_feature
        }
        
        print(f"  - Identity features initialized (Trump: {trump_feature.shape}, Cage: {cage_feature.shape})")
    
    def _simulate_face_swap(self, images, swap_type):
        """
        模拟人脸交换（用于测试和回退）
        当SimSwap不可用时使用
        
        Args:
            images: 输入图像
            swap_type: 交换类型
            
        Returns:
            模拟的人脸交换结果
        """
        # 创建模拟的人脸交换效果
        simulated = images.clone()
        
        # 根据交换类型添加不同的变换
        if swap_type in ['A', 'trump_to_cage']:
            # 模拟Trump->Cage的变换：稍微调暗图像，添加轻微噪声
            simulated = simulated * 0.95 + torch.randn_like(simulated) * 0.03
            
        elif swap_type in ['B', 'cage_to_trump']:
            # 模拟Cage->Trump的变换：稍微调亮图像，添加不同的噪声
            simulated = simulated * 1.05 + torch.randn_like(simulated) * 0.02
            
        # 确保输出在合理范围内
        if images.min() >= -0.1:  # 如果输入接近[-1,1]范围
            simulated = torch.clamp(simulated, -1, 1)
        else:  # 如果输入接近[0,1]范围
            simulated = torch.clamp(simulated, 0, 1)
        
        return simulated
    
    def batch_deepfake(self, encoded_images, swap_types):
        """
        批量人脸交换（支持不同的交换类型）
        
        Args:
            encoded_images: 输入图像批次
            swap_types: 交换类型列表，与批次大小相同
            
        Returns:
            批量处理后的图像
        """
        batch_size = encoded_images.shape[0]
        results = []
        
        for i in range(batch_size):
            single_image = encoded_images[i:i+1]  # 保持批次维度
            swap_type = swap_types[i] if isinstance(swap_types, list) else swap_types
            
            result = self.deepfake(single_image, swap_type)
            results.append(result)
        
        return torch.cat(results, dim=0)
    
    def preprocess_images(self, images):
        """
        预处理图像（如果需要特殊的预处理）
        
        Args:
            images: 输入图像
            
        Returns:
            预处理后的图像
        """
        # 确保图像在正确的范围内
        if images.max() > 1.1:  # 如果是[0,255]范围
            images = images / 255.0
        
        if images.min() >= 0 and images.max() <= 1:  # 如果是[0,1]范围
            images = images * 2.0 - 1.0  # 转换到[-1,1]
            
        return images
    
    def save_model_info(self, path):
        """保存模型信息"""
        info = {
            'is_available': self.is_available,
            'use_simulation': self.use_simulation,
            'device': str(self.device),
            'simswap_model_path': self.simswap_model_path
        }
        
        torch.save(info, path)
        print(f"SimSwap model info saved to {path}")

def test_simswap_gan_simulator():
    """测试SimSwap GAN模拟器"""
    print("=== Testing SimSwap GAN Simulator ===")
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # 创建SimSwap模拟器
    simulator = SimSwapGANSimulator(
        device=device,
        simswap_model_path=None  # 使用默认路径
    )
    
    # 测试数据
    batch_size = 4
    image_size = 256
    
    # 创建测试图像（模拟嵌入水印的图像）
    test_images = torch.randn(batch_size, 3, image_size, image_size)
    if device == 'cuda':
        test_images = test_images.cuda()
    
    print(f"Test images shape: {test_images.shape}")
    print(f"Test images range: [{test_images.min().item():.3f}, {test_images.max().item():.3f}]")
    
    # 测试单一交换类型
    print("\n--- Testing single swap type ---")
    
    # Trump -> Cage
    deepfaked_A = simulator.deepfake(test_images, swap_type='A')
    print(f"Trump->Cage result shape: {deepfaked_A.shape}")
    print(f"Trump->Cage result range: [{deepfaked_A.min().item():.3f}, {deepfaked_A.max().item():.3f}]")
    
    # Cage -> Trump
    deepfaked_B = simulator.deepfake(test_images, swap_type='B')
    print(f"Cage->Trump result shape: {deepfaked_B.shape}")
    print(f"Cage->Trump result range: [{deepfaked_B.min().item():.3f}, {deepfaked_B.max().item():.3f}]")
    
    # 测试批量交换
    print("\n--- Testing batch swap ---")
    swap_types = ['A', 'B', 'A', 'B']  # 交替交换
    batch_results = simulator.batch_deepfake(test_images, swap_types)
    print(f"Batch results shape: {batch_results.shape}")
    print(f"Batch results range: [{batch_results.min().item():.3f}, {batch_results.max().item():.3f}]")
    
    # 测试预处理
    print("\n--- Testing preprocessing ---")
    test_images_255 = torch.randint(0, 256, (2, 3, 64, 64), dtype=torch.float32)
    if device == 'cuda':
        test_images_255 = test_images_255.cuda()
    
    preprocessed = simulator.preprocess_images(test_images_255)
    print(f"Preprocessed from [0,255] to [{preprocessed.min().item():.3f}, {preprocessed.max().item():.3f}]")
    
    print("SimSwap GAN Simulator test completed!\n")

if __name__ == "__main__":
    test_simswap_gan_simulator() 